package com.nulldev.util.security.encryptAPI.impls.aes128_counter;

import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import com.nulldev.util.VariableAPI.ByteUtil;
import com.nulldev.util.data.Base32.Base32;
import com.nulldev.util.data.Base64.Base64;
import com.nulldev.util.security.encryptAPI.impls.aes128.KeyEncoding;

public class AESCipherCounter {
	private static final String ALGORITHM_AES256 = "AES/CTR/NoPadding";
	// ECP, default
//    private static final String ALGORITHM_AES256 = "AES";
	private static final byte[] INITIAL_IV =
		{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
	private final SecretKeySpec secretKeySpec;
	private final Cipher cipher;
	private IvParameterSpec iv;

	/**
	 * Create AESCipher based on existing {@link Key}
	 *
	 * @param key Key
	 */
	public AESCipherCounter(Key key) {
		this(key.getEncoded());
	}

	/**
	 * Create AESCipher based on existing {@link Key} and Initial Vector (iv) in
	 * bytes
	 *
	 * @param key Key
	 */
	public AESCipherCounter(Key key, byte[] iv) {
		this(key.getEncoded(), iv);
	}

	/**
	 * <p>
	 * Create AESCipher using a byte[] array as a key
	 * </p>
	 * <p/>
	 * <p>
	 * <strong>NOTE:</strong> Uses an Initial Vector of 16 0x0 bytes. This should
	 * not be used to create strong security.
	 * </p>
	 *
	 * @param key Key
	 */
	public AESCipherCounter(byte[] key) {
		this(key, INITIAL_IV);
	}

	private AESCipherCounter(byte[] key, byte[] iv) {
		try {
			this.secretKeySpec = new SecretKeySpec(key, "AES");
			this.iv = new IvParameterSpec(iv);
			this.cipher = Cipher.getInstance(ALGORITHM_AES256);
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * Takes message and encrypts with Key
	 *
	 * @param message String
	 * @return String Base64 encoded
	 */
	public String getEncryptedMessage(String message) {
		try {
			Cipher cipher = getCipher(Cipher.ENCRYPT_MODE);

			byte[] encryptedTextBytes = cipher.doFinal(message.getBytes("UTF-8"));

			return Base64.encodeBase64String(encryptedTextBytes);
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * Takes message and encrypts with Key
	 *
	 * @param message String
	 * @return String Base64 encoded
	 */
	public String getEncryptedMessageHEX(String message) {
		try {
			Cipher cipher = getCipher(Cipher.ENCRYPT_MODE);

			byte[] encryptedTextBytes = cipher.doFinal(message.getBytes("UTF-8"));

			return ByteUtil.getHex(encryptedTextBytes);
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * Takes message and encrypts with Key
	 *
	 * @param message bytes
	 * @return encrypted bytes
	 */
	public byte[] getEncryptedMessage(final byte[] message) {
		try {
			Cipher cipher = getCipher(Cipher.ENCRYPT_MODE);

			byte[] encryptedTextBytes = cipher.doFinal(message);

			return encryptedTextBytes;
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * Takes Base64 encoded String and decodes with provided key
	 *
	 * @param message String encoded with Base64
	 * @return String
	 */
	public String getDecryptedMessage(String message) {
		try {
			Cipher cipher = getCipher(Cipher.DECRYPT_MODE);

			byte[] encryptedTextBytes = Base64.decodeBase64(message);
			byte[] decryptedTextBytes = cipher.doFinal(encryptedTextBytes);

			return new String(decryptedTextBytes);
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * Takes Base64 encoded String and decodes with provided key
	 *
	 * @param message String encoded with Base64
	 * @return String
	 */
	public String getDecryptedMessageHEX(String message) {
		try {
			Cipher cipher = getCipher(Cipher.DECRYPT_MODE);

			byte[] encryptedTextBytes = ByteUtil.fromHex(message);
			byte[] decryptedTextBytes = cipher.doFinal(encryptedTextBytes);

			return new String(decryptedTextBytes);
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * Takes Base64 encoded String and decodes with provided key
	 *
	 * @param message bytes
	 * @return decrypted bytes
	 */
	public byte[] getDecryptedMessage(final byte[] message) {
		try {
			Cipher cipher = getCipher(Cipher.DECRYPT_MODE);

			byte[] decryptedTextBytes = cipher.doFinal(message);

			return decryptedTextBytes;
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * Get IV in Base64 Encoded String
	 *
	 * @return String Base64 Encoded
	 */
	public String getIV() {
		return Base64.encodeBase64String(iv.getIV());
	}

	/**
	 * Base64 encoded version of key
	 *
	 * @return String
	 */
	public String getKey() {
		return getKey(KeyEncoding.Base64);
	}

	public String getKey(KeyEncoding encoding) {
		String result = null;
		switch (encoding) {
			case Base64:
				result = Base64.encodeBase64String(secretKeySpec.getEncoded());
				break;
			case Hex:
				result = ByteUtil.getHex(secretKeySpec.getEncoded());
				break;
			case Base32:
				result = Base32.encode(secretKeySpec.getEncoded());
				break;
		}

		return result;
	}

	private Cipher getCipher(int encryptMode) throws InvalidKeyException, InvalidAlgorithmParameterException {
		cipher.init(encryptMode, getSecretKeySpec(), iv);
		return cipher;
	}

	private SecretKeySpec getSecretKeySpec() {
		return secretKeySpec;
	}
}
